package health_system.health_system.websocket.Test;

import jakarta.websocket.*;
import jakarta.websocket.server.ServerEndpoint;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.server.standard.SpringConfigurator;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ServerEndpoint(value = "/upload", configurator = SpringConfigurator.class) // 添加 configurator

public class FileUploadWebSocket {

    private ConcurrentHashMap<String, ByteArrayOutputStream> fileBuffers = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String, Session> sessions = new ConcurrentHashMap<>();

    @Value("${upload.dir:./uploads}")
    private String uploadDir;

    @OnOpen
    public void onOpen(Session session) {
        session.setMaxIdleTimeout(300000); // 5分钟超时
        session.setMaxBinaryMessageBufferSize(10485760); // 10MB 缓冲区
        sessions.put(session.getId(), session);
        System.out.println("客户端连接: " + session.getId());
        checkUploadDirectory();
    }
    private boolean isAllowedClient(String clientIp) {
        return true;
    }
    @OnClose
    public void onClose(Session session) {
        sessions.remove(session.getId());
        System.out.println("客户端断开: " + session.getId());
    }

    @OnMessage
    public void handleBinaryMessage(Session session, ByteBuffer byteBuffer, boolean isLast) {
        try {
            String sessionId = session.getId();
            ByteArrayOutputStream bufferStream = fileBuffers.computeIfAbsent(sessionId, k -> new ByteArrayOutputStream());

            // 写入当前分片数据
            byte[] bytes = new byte[byteBuffer.remaining()];
            byteBuffer.get(bytes);
            bufferStream.write(bytes);

            if (isLast) { // 所有分片接收完成
                byte[] fullData = bufferStream.toByteArray();
                fileBuffers.remove(sessionId); // 清理缓存

                String fileName = generateUniqueFileName();
                Path filePath = Paths.get(uploadDir, fileName);
                Files.createDirectories(filePath.getParent());
                Files.write(filePath, fullData);

                session.getAsyncRemote().sendText("图片已保存为: " + fileName); // 异步发送响应
                System.out.println("文件保存成功: " + fileName);
            }
        } catch (IOException e) {
            handleError(session, "文件保存失败: " + e.getMessage());
            fileBuffers.remove(session.getId()); // 确保异常时清理缓存
        }
    }

    @OnError
    public void onError(Session session, Throwable throwable) {
        System.err.println("WebSocket错误: " + throwable.getMessage());
        throwable.printStackTrace(); // 打印完整堆栈
        handleError(session, "服务器错误: " + throwable.getMessage());
    }
    private String generateUniqueFileName() {
        DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyyMMddHHmmssSSS");
        String timestamp = LocalDateTime.now().format(formatter);
        int random = (int) (Math.random() * 1000);
        return String.format("image_%s_%03d.jpg", timestamp, random);
    }

    private void checkUploadDirectory() {
        System.out.println("检查上传目录，路径: " + uploadDir); // 添加日志
        try {
            Path path = Paths.get(uploadDir);
            if (!Files.exists(path)) {
                Files.createDirectories(path);
                System.out.println("创建上传目录: " + path.toAbsolutePath());
            }
        } catch (IOException e) {
            System.err.println("无法创建上传目录: " + e.getMessage());
        }
    }

    private void handleError(Session session, String errorMsg) {
        try {
            if (session.isOpen()) {
                session.getBasicRemote().sendText("ERROR: " + errorMsg);
            }
        } catch (IOException e) {
            System.err.println("发送错误消息失败: " + e.getMessage());
        }
    }
}